# -*- coding: utf-8 -*-
# @Time    : 2021/11/13 21:00
# @Author  : travel2040
# @Email   : travel2040@163.com
# @File    : baidu-aip-百度ai接口读取视频定位人体并截图.py
# 需要安装的库
# pip install baidu-aip
# pip install opencv_python

import cv2
import time
import os
import random
from datetime import datetime
from aip import AipBodyAnalysis

class ClipStream():
    def __init__(self,min_score):

        APP_ID = '444444433'
        API_KEY = 'jpVq4v2PmG60iFBySVE80oQ9'
        SECRET_KEY = 'iSZSKNrBCZDAvLD8LXQ1tHrlhW55AA2p'

        self.client = AipBodyAnalysis(APP_ID, API_KEY, SECRET_KEY)
        self.min_score = min_score


    def check(self,frame,result):

        # point_color = (0, 255, 0)  # BGR

        # person_num = result["person_num"]

        person_info = result["person_info"]

        frame_full_name = "%s_%d"%(datetime.now().strftime('%Y%m%d_%H%M%S'),random.randint(1000,9999))
        frame_y = False

        for person in person_info:
            location = person["location"]
            l_left = int(location["left"])
            l_top = int(location["top"])
            l_width = int(location["width"])
            l_height = int(location["height"])
            l_score = location["score"]
            # body_parts = person["body_parts"]
            attributes = person["attributes"]

            if l_width>100 and l_height>280 and l_score > self.min_score : #l en(body_parts.keys())>=12
            # if l_score > self.min_score:
                # print(attributes)

                orientation = attributes["orientation"]
                orientation_name = orientation["name"] # 正面 or 背面 or 左侧面 or 右侧面
                gender = attributes["gender"]
                gender_name = gender["name"] # 男性 or 女性

                occlusion = attributes["occlusion"]
                occlusion_name = occlusion["name"] # 无遮挡 重度遮挡  轻度遮挡

                # 人体属性
                # {'attributes': {'is_human': {'name': '正常人体', 'score': 0.9839658737182617}, 'upper_cut': {'name': '无上方截断', 'score': 0.9930918216705322}, 'occlusion': {'name': '重度遮挡', 'score': 0.9599432349205017}, 'age': {'name': '青年', 'score': 0.9877013564109802}, 'carrying_item': {'name': '不确定', 'score': 0.6675317883491516}, 'umbrella': {'name': '未打伞', 'score': 0.9997518658638}, 'lower_cut': {'name': '有下方截断', 'score': 0.6453482508659363}, 'bag': {'name': '无背包', 'score': 0.8834974765777588}, 'orientation': {'name': '背面', 'score': 0.953889012336731}, 'vehicle': {'name': '无交通工具', 'score': 0.9997878670692444}, 'glasses': {'name': '不确定', 'score': 0.9762601256370544}, 'cellphone': {'name': '不确定', 'score': 0.9999989867210388}, 'headwear': {'name': '无帽', 'score': 0.9997523427009583}, 'upper_wear_fg': {'name': '外套', 'score': 0.4856634438037872}, 'face_mask': {'name': '不确定', 'score': 0.6241160035133362}, 'lower_wear': {'name': '不确定', 'score': 0.9736524820327759}, 'gender': {'name': '男性', 'score': 0.6472302079200745}, 'upper_wear': {'name': '长袖', 'score': 0.5240340232849121}, 'smoke': {'name': '不确定', 'score': 0.94351726770401}, 'lower_color': {'name': '不确定', 'score': 0.9787816405296326}, 'upper_color': {'name': '白', 'score': 0.9357854127883911}, 'upper_wear_texture': {'name': '纯色', 'score': 0.9872881174087524}}, 'location': {'height': 137, 'width': 73, 'left': 868, 'top': 77, 'score': 0.1184862479567528}}

                if  orientation_name == "正面" and  occlusion_name != "重度遮挡":
                # if True:
                    frame_y = True
                    print("发现目标")
                    # 生成截图
                    crop_name = "%s_%.2f_%s_%s_%s.jpg"%(frame_full_name,l_score,gender_name,occlusion_name,orientation_name)

                    # 保存截图
                    crop_filename = os.path.join(self.clip_dir,crop_name)
                    crop_image = frame[l_top:l_top + l_height, l_left:l_left + l_width]
                    cv2.imencode('.jpg', crop_image)[1].tofile(crop_filename)


                    # cv2.rectangle(frame, (l_left, l_top), (l_left+l_width, l_top+l_height), (0, 255, 0), 1)
                    # cv2.rectangle(img, 左上角，右下角，color，宽度)


                    # cv2.imshow("clip_image", clip_image)
                    # cv2.waitKey(0)
                    # cv2.destroyAllWindows()

                    # 渲染frame
                    # pt1 = (l_left,l_top)
                    # pt2 = (l_left+l_width,l_top+l_height)
                    # thickness = 1
                    # lineType = 8
                    # cv2.rectangle(frame,pt1,pt2,point_color,thickness,lineType)

        if frame_y:
            # 保存完整图片
            ff = os.path.join(self.clip_dir, frame_full_name+".jpg")
            cv2.imencode('.jpg', frame)[1].tofile(ff)

    def run(self,address,prefix,need_frame,skip_frame=25):
        print("开始执行>>>",address)
        self.clip_dir = "D:\\file\\camera_v1"
        # self.clip_dir = "/home/sun/Documents/file/camera"


        self.clip_dir = os.path.join(self.clip_dir, prefix)
        # self.clip_dir = os.path.join(self.clip_dir, datetime.now().strftime('%Y%m%d'))

        if not os.path.exists(self.clip_dir):
            os.makedirs(self.clip_dir)

        cap = cv2.VideoCapture(address)
        try:
            num = 0
            handle_num = 0
            while True:
                num += 1
                ret, frame = cap.read()
                if ret:
                    # cv2.imshow("frame", frame)

                    if num % skip_frame ==0:
                        time.sleep(1.5)
                        handle_num +=1
                        # time_start = time.time()
                        encoded_ret, encoded_image = cv2.imencode(".jpg", frame)
                        image = encoded_image.tobytes()
                        # result = self.client.bodyAnalysis(image)
                        result = self.client.bodyAttr(image)
                        self.check(frame, result)
                        # print("frame speed", time.time() - time_start)

                        # cv2.imshow("frame", frame)

                        if handle_num > need_frame:
                            break
                    if cv2.waitKey(20) & 0xFF == ord('q'):
                        break
                else:
                    break
        except Exception as e:
            print(e)
            cap.release()
            cv2.destroyAllWindows()

if __name__ == '__main__':

    clipStream = ClipStream(min_score=0.2)
    clipStream.run("rtsp://admin:a12345678@192.168.0.1", prefix="crop", need_frame=100, skip_frame=40)
